from pymilvus import (
    connections,
    Collection,
)
from zhipuai import ZhipuAI
import os
from dotenv import load_dotenv, find_dotenv


def query_milvus_by_text(query_text):
    """
    根据给定的文本查询Milvus集合中相似的记录。
    
    参数:
    query_text (str): 用于查询的文本。
    """
    # 加载环境变量
    load_dotenv(find_dotenv())
    api_key = os.getenv('ZHIPU_API_KEY')

    # 初始化ZhipuAI客户端
    client = ZhipuAI(api_key=api_key)

    # 确保连接到Milvus
    connections.connect("default", host='127.0.0.1', port='19530')

    # 获取集合实例
    collection = Collection("text_collection")

    # 加载集合到内存
    collection.load()

    # 检查并创建索引（如果尚未创建）
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128},
    }
    if not collection.has_index():
        collection.create_index(field_name="vector", index_params=index_params)

    # 将查询文本转换为向量
    response = client.embeddings.create(
        model="embedding-2",
        input=query_text,
    )
    query_vector = response.data[0].embedding

    # 构造查询参数
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10},
    }

    # 执行查询
    results = collection.search(
        data=[query_vector],
        anns_field="vector",
        param=search_params,
        limit=10,
        output_fields=["text"],
    )

    # 初始化最小距离为无穷大
    min_distance = float('inf')
    closest_text = ""

    for hits in results:
        for hit in hits:
            print(f" - Text: {hit.entity.get('text')}, Distance: {hit.distance}")

    # 打印查询结果
    for hits in results:
        for hit in hits:
            distance = hit.distance
            if distance < min_distance:
                min_distance = distance
                closest_text = hit.entity.get('text')

    return closest_text, min_distance


if __name__ == "__main__":
    # 使用方法示例
    query_text_example = "2019"
    closest_text, min_distance = query_milvus_by_text(query_text_example)
    print(f"The text most similar to '{query_text_example}' is: '{closest_text}', with a distance of {min_distance}.")
